feat(kernel): support federated token providers - #444
Conversation
|
For Go, we should also reuse-existing FederationProvider to avoid confusion as much as possible. Can we thread that through? Long term is support custom token provider in kernel which is out of scope for now |
vuanhphung
changed the title
![peco-review-bot[bot]](https://avatars.githubusercontent.com/in/3801177?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: 2 Low
Looks good — the federated-provider kernel wiring is correct: the wrapper preserves the Thrift path, GetToken is snapshotted exactly once (telemetry short-circuit verified by test), and the federation client ID is correctly gated to the PAT branch. Two minor low-severity items: one stale retry-range comment the kernel bump left behind, and missing coverage for the two new federated error paths.
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the kernel federated-token path is implemented correctly: the base provider is snapshotted once (kernel does its own exchange via set_identity_federation_client_id), the telemetry classifier short-circuits to avoid a second GetToken, and the new federatedTokenAuthenticator wrapper is safe on the Thrift path (nothing type-asserts the concrete provider, and the embedded authenticator implements neither M2M/U2M interface). One low doc-consistency gap: CONNECTION_PARAMETERS.md's kernel notes weren't updated with the snapshot/expiry caveat that README got.
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean reuse of WithFederatedTokenProvider* on the kernel path with correct type-switch ordering, context-threaded token resolution, and a telemetry short-circuit that avoids a second provider snapshot. One low-severity question about the account-wide (no client-id) federation path.
|
@eric-wang-1990 PTAL! |
0d6a9eb to
0dbbf2b
Compare
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean reuse of the federation authenticator that leaves the Thrift path byte-identical (embedded delegate) while snapshotting one token for the kernel, with correct type-switch ordering and no double GetToken (telemetry short-circuits). One low note: the kernel snapshots the base provider rather than the FederationProvider, so the account-wide (no-client-id) case relies on the kernel doing its own exchange — worth confirming that path is e2e-tested.
| // Wrap with federation provider that auto-detects need for token exchange | ||
| federationProvider := tokenprovider.NewFederationProvider(baseProvider, c.Host) | ||
| c.Authenticator = tokenprovider.NewAuthenticator(federationProvider) | ||
| c.Authenticator = &federatedTokenAuthenticator{ |
There was a problem hiding this comment.
Does this change affect existing Go Thrift side token-exchange feature?
There was a problem hiding this comment.
@eric-wang-1990 nope, this is just adding more metadata info. You can think of it like this
Before:
c.Authenticator = oldAuthenticator
After:
c.Authenticator = wrapper{
Authenticator: oldAuthenticator,
// kernel-only metadata
}
calling c.Authenticator.Authenticate() still delegate to oldAuthenticator
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean, well-tested addition of federated-token support on the kernel backend. The base provider is snapshotted exactly once (telemetry short-circuits federatedTokenAuthenticator to avoid a second GetToken), context is properly threaded through validateKernelConfigContext/resolveKernelAuthContext, the Thrift exchange path is preserved and covered by a new httptest, and docs/CHANGELOG are updated consistently. One low-severity design question about the account-wide path is noted inline.
| if token == nil || token.AccessToken == "" { | ||
| return kernel.Auth{}, errors.New("databricks: the federated token provider returned an empty token") | ||
| } | ||
| return kernel.Auth{Mode: kernel.AuthPAT, Token: token.AccessToken, ClientID: a.clientID}, nil |
There was a problem hiding this comment.
🔵 Low — For the account-wide case (WithFederatedTokenProvider, no client ID), a.clientID is empty, so setAuth skips kernel_session_config_set_identity_federation_client_id and hands the raw external-IdP subject token to the kernel via set_auth_pat with no federation signal at all. The exchange the Thrift path performs in-driver (FederationProvider.GetToken) is not done here — it's delegated to the kernel, which for the SP-wide case is cued by the federation client ID but for account-wide has only the bare token to auto-detect from. Worth confirming that make test-kernel actually exercised the account-wide federated path (not just skipped it credential-gated) so the kernel is verified to auto-detect federation from the token alone; otherwise a raw external token could be sent straight through as a bearer PAT and rejected. If the kernel's account-wide auto-detection is a known guarantee, this is fine as-is.
(Anchored to the nearest changed line — see the description for the exact location.)
2 participants
Reuses
WithFederatedTokenProvider*on the kernel backend. The driver snapshots one base-provider token for kernel-side PAT federation;WithFederatedTokenProviderAndClientIDalso forwards the SP-wide client ID. Because the kernel cannot refresh the provider, an expired token requires a new connection.Pins
KERNEL_REVto the federation-capable kernel. The kernel bump also requires aligning one retry-setter assertion with its clamp behavior.Validation:
go test ./...,make test-kernel, andmake lint.